import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.pqc.crypto.crystals.dilithium.*;
import org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider;
import org.bouncycastle.pqc.jcajce.spec.DilithiumParameterSpec;
// 1. 加上缺失的 import
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.pqc.crypto.crystals.dilithium.*;
import org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.util.io.pem.PemObject;
import org.bouncycastle.util.io.pem.PemWriter;
import org.bouncycastle.pqc.jcajce.provider.dilithium.BCDilithiumPublicKey;
import org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider;
import java.io.*;
import java.math.BigInteger;
import java.security.Security;
import java.security.SecureRandom;
import java.security.cert.X509Certificate;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Base64;
import java.util.Date;
import java.io.*;
import org.bouncycastle.asn1.*;
import org.bouncycastle.asn1.ocsp.*;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.*;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.cert.ocsp.*;
import org.bouncycastle.cert.ocsp.jcajce.JcaCertificateID;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.OperatorCreationException;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.bouncycastle.operator.jcajce.JcaDigestCalculatorProviderBuilder;

import java.io.*;
import java.math.BigInteger;
import java.security.PrivateKey;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Date;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.cert.ocsp.*;
import org.bouncycastle.cert.ocsp.jcajce.JcaCertificateID;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.bouncycastle.operator.jcajce.JcaDigestCalculatorProviderBuilder;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;

import java.io.*;
import java.math.BigInteger;
import java.security.PrivateKey;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Date;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.cert.ocsp.*;
import org.bouncycastle.cert.ocsp.jcajce.JcaCertificateID;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.bouncycastle.operator.jcajce.JcaDigestCalculatorProviderBuilder;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;

import java.io.*;
import java.math.BigInteger;
import java.security.PrivateKey;
import java.security.Security;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Date;
// 2. 构造 AlgorithmIdentifier 时把 String 转成 ASN1ObjectIdentifier
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import java.math.BigInteger;
import java.security.*;
import java.security.cert.X509Certificate;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Date;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import org.bouncycastle.pqc.crypto.crystals.dilithium.*;
import org.bouncycastle.util.encoders.Hex;
import java.security.SecureRandom;

public class JcaDilithiumCertExample {
    static {
        Security.addProvider(new BouncyCastleProvider());   // 先加标准 BC
        Security.addProvider(new BouncyCastlePQCProvider()); // 再加 PQC 扩展
    }
    public static void main(String[] args) throws Exception {

        // 2. 生成密钥对
        SecureRandom random = new SecureRandom();

        // 1. 选择安全参数（Dilithium2, Dilithium3, Dilithium5）
        DilithiumParameters params = DilithiumParameters.dilithium5;
        DilithiumKeyPairGenerator keyPairGen = new DilithiumKeyPairGenerator();
        keyPairGen.init(new DilithiumKeyGenerationParameters(random, params));
        var keyPair = keyPairGen.generateKeyPair();

        String publicoid = "2.16.840.1.101.3.4.3.17";

        //IDs: { 2.16.840.1.101.3.4.3.17, id-ml-dsa-44, ML-DSA-44, MLDSA44 } @ default
        if(params == DilithiumParameters.dilithium2){
            publicoid = "2.16.840.1.101.3.4.3.17";
        }
        //IDs: { 2.16.840.1.101.3.4.3.18, id-ml-dsa-65, ML-DSA-65, MLDSA65 } @ default
        if(params == DilithiumParameters.dilithium3){
            publicoid = "2.16.840.1.101.3.4.3.18";
        }
        //IDs: { 2.16.840.1.101.3.4.3.19, id-ml-dsa-87, ML-DSA-87, MLDSA87 } @ default
        if(params == DilithiumParameters.dilithium5){
            publicoid = "2.16.840.1.101.3.4.3.19";
        }
        DilithiumPublicKeyParameters pubKey = (DilithiumPublicKeyParameters) keyPair.getPublic();
        DilithiumPrivateKeyParameters privKey = (DilithiumPrivateKeyParameters) keyPair.getPrivate();

        //PublicKey  pub  = kp.getPublic();
        //PrivateKey priv = kp.getPrivate();

        /* 2. 证书内容 */
        X500Name subject = new X500Name("CN=Test Dilithium Self-Signed, O=Demo, C=CN");
        BigInteger serial = BigInteger.valueOf(System.currentTimeMillis());
        Instant now = Instant.now();
        Date notBefore = Date.from(now);
        Date notAfter  = Date.from(now.plus(365, ChronoUnit.DAYS));

        /* 3. 签名算法 OID -> ML-DSA-44 */
AlgorithmIdentifier sigAlgId =
        new AlgorithmIdentifier(new ASN1ObjectIdentifier(publicoid));

        AlgorithmIdentifier keyAlgId = new AlgorithmIdentifier(
            new ASN1ObjectIdentifier(publicoid)); //弃用 Dilithium OID 1.3.6.1.4.1.2.267.1.6.5
        /* 4. 构造证书 */
        SubjectPublicKeyInfo pubInfo =new SubjectPublicKeyInfo(keyAlgId, pubKey.getEncoded());
                //SubjectPublicKeyInfo.getInstance(pubKey.getEncoded());
        X509v3CertificateBuilder certGen =
                new X509v3CertificateBuilder(subject, serial, notBefore, notAfter,
                                             subject, pubInfo);

        /* 5. 签名器 */
        ContentSigner signer = new ContentSigner() {
            private ByteArrayOutputStream buf = new ByteArrayOutputStream();
            @Override public AlgorithmIdentifier getAlgorithmIdentifier() { return sigAlgId; }
            @Override public OutputStream getOutputStream() { return buf; }
            @Override public byte[] getSignature() {
           
                DilithiumSigner signer = new DilithiumSigner();
                signer.init(true, privKey); // true 表示用于签名
                byte[] signature = signer.generateSignature(buf.toByteArray());
                return signature;
 
                    }
                };

        /* 6. 生成证书 */
        X509Certificate cert = new JcaX509CertificateConverter()
                .setProvider("BC")
                .getCertificate(certGen.build(signer));

        /* 7. 验证 & 打印 */

// 1. 把 BC-PQC 参数对象包装成 JCA PublicKey
BCDilithiumPublicKey jcaPub = new BCDilithiumPublicKey(pubKey);

// 2. 再验证
//cert.verify(jcaPub);
        System.out.println("证书生成成功，Subject: " + cert.getSubjectX500Principal().getName());
        System.out.println("签名算法: " + cert.getSigAlgOID());
        // 1. 将证书写出为 PEM 格式
        writeCertificateToPEM(cert, "dilithium_cert.crt");
        
        // 2. 将私钥写出为 PEM 格式
        writePrivateKeyToPEM(privKey, "dilithium_private_key.pem");
        
        // 3. 将公钥写出为 PEM 格式（可选）
        writePublicKeyToPEM(pubKey, "dilithium_public_key.pem");
        
        // 4. 将证书和私钥写出为 DER 格式（二进制）
        writeCertificateToDER(cert, "dilithium_cert.cer");
        writePrivateKeyToDER(privKey, "dilithium_private_key.der");
        
        System.out.println("所有文件已成功写出！");
    }

    // 将证书写出为 PEM 格式
    private static void writeCertificateToPEM(X509Certificate cert, String filename) throws Exception {
        try (PemWriter pemWriter = new PemWriter(new FileWriter(filename))) {
            PemObject pemObject = new PemObject("CERTIFICATE", cert.getEncoded());
            pemWriter.writeObject(pemObject);
        }
        System.out.println("证书已保存为 PEM 格式: " + filename);
    }

    // 将私钥写出为 PEM 格式
    private static void writePrivateKeyToPEM(DilithiumPrivateKeyParameters privKey, String filename) throws Exception {
        try (PemWriter pemWriter = new PemWriter(new FileWriter(filename))) {
            PemObject pemObject = new PemObject("PRIVATE KEY", privKey.getEncoded());
            pemWriter.writeObject(pemObject);
        }
        System.out.println("私钥已保存为 PEM 格式: " + filename);
    }

    // 将公钥写出为 PEM 格式
    private static void writePublicKeyToPEM(DilithiumPublicKeyParameters pubKey, String filename) throws Exception {
        try (PemWriter pemWriter = new PemWriter(new FileWriter(filename))) {
            PemObject pemObject = new PemObject("PUBLIC KEY", pubKey.getEncoded());
            pemWriter.writeObject(pemObject);
        }
        System.out.println("公钥已保存为 PEM 格式: " + filename);
    }

    // 将证书写出为 DER 格式（二进制）
    private static void writeCertificateToDER(X509Certificate cert, String filename) throws Exception {
        try (FileOutputStream fos = new FileOutputStream(filename)) {
            fos.write(cert.getEncoded());
        }
        System.out.println("证书已保存为 DER 格式: " + filename);
    }

    // 将私钥写出为 DER 格式（二进制）
    private static void writePrivateKeyToDER(DilithiumPrivateKeyParameters privKey, String filename) throws Exception {
        try (FileOutputStream fos = new FileOutputStream(filename)) {
            fos.write(privKey.getEncoded());
        }
        System.out.println("私钥已保存为 DER 格式: " + filename);
    }

    // 可选：生成 PKCS#12 文件（包含证书和私钥）
    private static void createPKCS12File(X509Certificate cert, DilithiumPrivateKeyParameters privKey, 
                                        String filename, String password) throws Exception {
        // 注意：由于 Dilithium 是后量子算法，标准的 PKCS#12 可能不完全支持
        // 这里提供一个概念性的实现
        System.out.println("PKCS#12 文件创建可能需要额外的配置来支持 Dilithium 算法");
        
        // 可以使用 BouncyCastle 的 PKCS12Builder 或其他工具
        // 具体实现取决于你的具体需求
    }

}